package com.yntx.usercenter.common.config.sharding;

import com.yntx.usercenter.base.cache.TenantCache;
import com.yntx.usercenter.base.mapper.TenantMapper;
import com.yntx.usercenter.base.po.Tenant;
import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
import org.apache.shardingsphere.infra.config.RuleConfiguration;
import org.apache.shardingsphere.infra.database.DefaultSchema;
import org.apache.shardingsphere.infra.datasource.props.DataSourceProperties;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.algorithm.config.AlgorithmProvidedShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.util.*;

/**
 * 项目启动后，读取已有分表，进行缓存
 *
 * @author wuzhengyu
 * @date 2023/05/07
 */
@Order(value = 1) // 数字越小，越先执行
@Component
public class ShardingTablesLoadRunner implements CommandLineRunner {

    @Autowired
    private TenantMapper tenantMapper;

    @Autowired
    private TenantCache tenantCache;

    @Autowired
    private ShardingSphereDataSource shardingSphereDataSource;

    @Override
    public void run(String... args) {
        // Sharding数据源更新
        updateShardingDataSource();

        // 读取已有分表，进行缓存
        ShardingAlgorithmTool.tableNameCacheReloadAll();
    }

    private void updateShardingDataSource() {
        List<Tenant> tenantList = tenantMapper.selectList(null);
        System.out.println(tenantList);

        ContextManager contextManager = shardingSphereDataSource.getContextManager();
        // 动态添加数据源
        Map<String, DataSourceProperties> dataSourcePropsMap = new HashMap<>();
        tenantList.forEach(tenant -> {
            Map<String, Object> props = new HashMap<>();
            props.put("type", "com.alibaba.druid.pool.DruidDataSource");
            props.put("driverClassName", tenant.getDatasourceDriver());
            props.put("url", tenant.getDatasourceUrl());
            props.put("username", tenant.getDatasourceUsername());
            props.put("password", tenant.getDatasourcePassword());
            DataSourceProperties dataSourceProperties = new DataSourceProperties("com.alibaba.druid.pool.DruidDataSource", props);
            dataSourcePropsMap.put(tenant.getTenantId(), dataSourceProperties);
            tenantCache.get().put(tenant.getTenantId(), tenant);
        });
        contextManager.alterDataSourceConfiguration(DefaultSchema.LOGIC_NAME, dataSourcePropsMap);

        // 分片规则重新设置
        Collection<RuleConfiguration> newRuleConfigList = new LinkedList<>();
        Collection<RuleConfiguration> oldRuleConfigList = contextManager.getMetaDataContexts()
                .getMetaData(DefaultSchema.LOGIC_NAME)
                .getRuleMetaData()
                .getConfigurations();

        for (RuleConfiguration oldRuleConfig : oldRuleConfigList) {
            if (oldRuleConfig instanceof AlgorithmProvidedShardingRuleConfiguration) {

                // Algorithm provided sharding rule configuration
                AlgorithmProvidedShardingRuleConfiguration oldAlgorithmConfig = (AlgorithmProvidedShardingRuleConfiguration) oldRuleConfig;
                AlgorithmProvidedShardingRuleConfiguration newAlgorithmConfig = new AlgorithmProvidedShardingRuleConfiguration();

                // Sharding table rule configuration Collection
                Collection<ShardingTableRuleConfiguration> newTableRuleConfigList = new LinkedList<>();
                Collection<ShardingTableRuleConfiguration> oldTableRuleConfigList = oldAlgorithmConfig.getTables();

                oldTableRuleConfigList.forEach(oldTableRuleConfig -> {
                    if ("t_user".equals(oldTableRuleConfig.getLogicTable())) {
                        ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(oldTableRuleConfig.getLogicTable(), oldTableRuleConfig.getActualDataNodes());
                        newTableRuleConfig.setTableShardingStrategy(oldTableRuleConfig.getTableShardingStrategy());
                        newTableRuleConfig.setDatabaseShardingStrategy(oldTableRuleConfig.getDatabaseShardingStrategy());
                        newTableRuleConfig.setKeyGenerateStrategy(oldTableRuleConfig.getKeyGenerateStrategy());

                        newTableRuleConfigList.add(newTableRuleConfig);
                    } else {
                        newTableRuleConfigList.add(oldTableRuleConfig);
                    }
                });

                newAlgorithmConfig.setTables(newTableRuleConfigList);
                newAlgorithmConfig.setAutoTables(oldAlgorithmConfig.getAutoTables());
                newAlgorithmConfig.setBindingTableGroups(oldAlgorithmConfig.getBindingTableGroups());
                newAlgorithmConfig.setBroadcastTables(oldAlgorithmConfig.getBroadcastTables());
                newAlgorithmConfig.setDefaultDatabaseShardingStrategy(oldAlgorithmConfig.getDefaultDatabaseShardingStrategy());
                newAlgorithmConfig.setDefaultTableShardingStrategy(oldAlgorithmConfig.getDefaultTableShardingStrategy());
                newAlgorithmConfig.setDefaultKeyGenerateStrategy(oldAlgorithmConfig.getDefaultKeyGenerateStrategy());
                newAlgorithmConfig.setDefaultShardingColumn(oldAlgorithmConfig.getDefaultShardingColumn());
                newAlgorithmConfig.setShardingAlgorithms(oldAlgorithmConfig.getShardingAlgorithms());
                newAlgorithmConfig.setKeyGenerators(oldAlgorithmConfig.getKeyGenerators());

                newRuleConfigList.add(newAlgorithmConfig);
            }
        }

        // update context
        contextManager.alterRuleConfiguration(DefaultSchema.LOGIC_NAME, newRuleConfigList);
    }

}

